'''
Author: Q.C.Li
Date: 2021-01-18 20:17:20
LastEditors: Q.C.Li
LastEditTime: 2021-05-21 16:12:33
Description: 使用点扩散函数计算小孔
'''
# %%
from scipy import signal
from scipy.fftpack import fftshift, ifftshift,fft,ifft

from CMD.calculate import *
from CMD.plotInitialize import *

imshow_ = lambda img,x,y,n: px.imshow(
    abs(img[::n,::n]),
    x=x[::n],y=y[::n]
    ).update_layout(yaxis_autorange=True)

# %% 计算小孔后的交叉谱密度
# 待调参数
sigma = 68e-6       # 光斑半径
factor = 1
xi = 62e-6    # 相干半径

L1 = 5e-3
d = 3e-3       # 小孔直径

q = xi/sigma
zeta = q/ (q**2+4)**.5
kappa = ((1-zeta)/(1+zeta))
n = 7
# 
'''
ipcluster start -n 2
'''
demo = PartialCoherentSR([sigma,sigma],[xi,xi],91.8,1)

# %% 3. 70m 处
L2 = 16e-3
# 传播距离
z1 = 25
z2 = 45

c1_old = L1*demo.k/(2*c.pi*z2)
c1 = (c.pi*z2)/((L1+L2)*np.sqrt(2)*demo.k)
c2 = d*demo.k/(2*c.pi*z2)
num = int(L2/c1)
# CSD 的坐标
xx_ctr = np.arange(-num*c1/2,num*c1/2,c1)
# psf 坐标
xx_mgn = np.arange(-num*c1,num*c1,c1)

psf = (d**2*np.sinc(c2*xx_mgn).reshape(-1,1) * np.sinc(c2*xx_mgn)).astype(np.complex128)

CSD = demo.GSM(z1+z2,xx_ctr,0)['CSD'].astype(np.complex128)

# %% 卷积
e_c = exp(1j*demo.k/(2*z2) * (
    (xx_ctr**2).reshape(-1,1) - xx_ctr**2) ).astype(np.complex128)
CSD2 = signal.convolve(CSD*e_c, psf,
    mode='valid',method='fft')


# %% 绘图需要的数据
xx_plot = xx_mgn[xx_ctr.shape[0]//2:-xx_ctr.shape[0]//2+1]

S = CSD2.diagonal()
SDC = CSD2/abs(np.sqrt(S)*np.sqrt(S.reshape(-1,1)))
# SDC = SDC/abs(SDC).max()

S0 = CSD.diagonal()
SDC0 = CSD/abs(S0[874]*S0.reshape(-1,1))**.5


# %% 绘图-有小孔-光强
name1 = f'孔:{d*1e3:.2f}mm'

fig1 = px.line(y=abs(S),x=xx_plot*1e3,
    title=name1,
    labels={'x':r'$x (mm)$','y':r'$I(x)$'})

# %% 有小孔-相干度
fig2 = px.line(y=abs(SDC[781]),x=2*xx_plot*1e3,title=name1,labels={'x':r'$\Delta x (mm)$','y':r'$\mu(\Delta x)$'})

# # %% 无小孔-相干度
# name2 = f'光斑半径/相干半径:{1/factor}'
# fig3 = px.line(y=abs(SDC0[::-1].diagonal()),x=2*xx_ctr*1e3,title=name2,labels={'x':r'$\Delta x (mm)$','y':r'$\mu(\Delta x)$'})

# # %% 无小孔-光强
# fig4 = px.line(y=abs(S0),x=xx_ctr*1e3,
#     title=name2,
#     labels={'x':r'$x (mm)$','y':r'$I(x)$'})


# %% subplolt-有小孔
figSP1 = make_subplots(2,1)

figSP1.add_traces([fig2.data[0],fig1.data[0]],[1,2],[1,1]);
figSP1.update_layout(
    title_text=name1,
    xaxis_title_text=r'$\Delta x (mm)$',
    yaxis_title_text=r'$\mu(\Delta x)$',
    xaxis2_title_text=r'$x (mm)$',
    yaxis2_title_text=r'$I(x)$' )


# %% subplot 无小孔
figSP2 = make_subplots(2,1)

figSP2.add_traces([fig3.data[0],fig4.data[0]],[1,2],[1,1]);
figSP2.update_layout(
    title_text=name2,
    xaxis_title_text=r'$\Delta x (mm)$',
    yaxis_title_text=r'$\mu(\Delta x)$',
    xaxis2_title_text=r'$x (mm)$',
    yaxis2_title_text=r'$I(x)$' )


# figSP2.write_image(rf"d:\Users\LIQIC\Desktop\小孔-相干性\sigma=1e-3\{1/factor:.1f}.png",scale=5)
px.line(abs(S))
px.line(abs(SDC[::1].diagonal()))
px.line(abs(CSD2[::-1].diagonal()))

px.imshow(abs(SDC))
px.imshow(abs(CSD2[::10,::10]))

